
#include <iostream>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>  
#include <direct.h> 
#include <vector>
#include "boost/filesystem/operations.hpp"
#include "boost/filesystem/path.hpp"
#include "boost/filesystem/operations.hpp"
#include "boost/filesystem/fstream.hpp"

#include "../txr_algo_dlm_rec_v1/common/define.h"
#include "../txr_algo_dlm_rec_v1/common/gpu_utility.h"
#include "yolo_v6.h"
#include <thread>
#include <chrono>


std::vector<std::string> EnumFiles(std::string path, std::string extension)
{
	namespace fs = boost::filesystem;
	fs::path _path(path);
	fs::path fullpath = fs::system_complete(_path);

	std::vector<std::string> ret;

	if (!fs::exists(fullpath) && !fs::is_directory(fullpath))
	{
		return ret;
	}
	fs::directory_iterator end_iter;
	for (fs::directory_iterator iter(fullpath); iter != end_iter; iter++)
	{
		try {
			if (fs::is_directory(*iter))
			{
				//ret.push_back(iter->path().wstring());
			}
			else if (iter->path().extension() == extension)
			{
				ret.push_back(iter->path().string());
			}
		}
		catch (const std::exception&)
		{
			continue;
		}
	}
	return ret;
}

void enum_models(std::string& mon_folder, std::string& dest_folder,std::string & mark_info)
{
	namespace fs = boost::filesystem;
	
	std::string gpu_model = GetGPUModel();
	if (gpu_model == "")
	{
		std::cout << "Unknown Gpu Model! \n";
		return;
	}
	//���Ŀ���ļ����Ƿ���Ч
	fs::path _dest_folder(dest_folder);
	if (!fs::exists(_dest_folder))
	{
		return;
	}
	std::vector<std::string> cfg_files = EnumFiles(mon_folder, std::string(".yaml"));
	for (int i = 0; i < cfg_files.size(); ++i)
	{
		st_trans_model_input input;
		fs::path _cfg_path(cfg_files[i]);
		input.cfg_path = _cfg_path.string();

		fs::path _model_path; 
		{
			fs::path _path = _cfg_path.replace_extension(".ennx");
			if (fs::exists(_path))
			{
				_model_path = _path;
			}
		} 
		{
			fs::path _path = _cfg_path.replace_extension(".onnx");
			if (fs::exists(_path))
			{
				_model_path = _path;
			}
		}
		if (_model_path.empty())
		{
			continue;
		}
		
		fs::path _gdats_path(dest_folder + "\\" + _cfg_path.stem().string() + mark_info + ".rec1");
		if (fs::exists(_gdats_path))
		{
			continue;
		}

		std::cout << "START===============================================\n";
		input.model_path = _model_path.string();
		input.gpu_model = gpu_model;
		input.trans_path = _gdats_path.string()+"-tmp";
		yolo_v6 c;
		if (c.TransModel(input))
		{
			std::cout << "DONE================================================\n";
			int result = std::rename(input.trans_path.c_str(), _gdats_path.string().c_str());
			if (result == 0) 
			{
				std::cout << "File renamed successfully.\n";
			} 
			else 
			{
				std::cout << "File rename failed.\n";
			}
		}
		else
		{
			std::cout << "FAIL!================================================\n";
		}
	}
}

int main(int argc, char* argv[])
{
	try
	{

		char cur_folder[260];
		_getcwd(cur_folder, sizeof(cur_folder));
		//printf("current working directory: \n %s\n", cur_folder);
		std::string cwd = cur_folder;

		std::string mon_folder = cwd;
		std::string dest_folder = cwd;
		std::string mark_info = "";
		if (argc >= 2)
		{
			mon_folder = argv[1];
			dest_folder = mon_folder;
		}
		if (argc >= 3)
		{
			dest_folder = argv[2];
		}
		if (argc >= 4)
		{
			mark_info = argv[3];
		}

		enum_models(mon_folder, dest_folder, mark_info);
		printf("done!\n");

	}
	catch (...)
	{
		printf("exception !\n");
	}

	for (int i = 0; i<5; ++i)
	{
		std::this_thread::sleep_for(std::chrono::milliseconds(1000));
		//printf("wait for shutdown %d !\n", 5 - i);
	}

	//getchar();
}
